from typing import List


class Solution:
    def removeBoxes(self, boxes: List[int]) -> int:
        size = len(boxes)

        # 定义状态矩阵：从i~j中移除颜色相同的盒子，j右边有k个和a[j]颜色相同的盒子
        dp = [[[0] * size for _ in range(size)] for _ in range(size)]

        for l in range(size):  # 遍历长度差
            for i in range(size - l):  # 遍历开始位置
                j = i + l  # 计算结束位置
                for k in range(0, size - j):
                    res = (k + 1) * (k + 1)
                    if j - 1 >= i:
                        res += dp[i][j - 1][0]
                    dp[i][j][k] = max(dp[i][j][k], res)

                    for t in range(i, j):
                        if boxes[t] == boxes[j]:
                            res = dp[i][t][k + 1]
                            if t + 1 <= j - 1:
                                res += dp[t + 1][j - 1][0]
                            dp[i][j][k] = max(dp[i][j][k], res)

        return dp[0][size - 1][0]


if __name__ == "__main__":
    print(Solution().removeBoxes([1, 3, 2, 2, 2, 3, 4, 3, 1]))  # 23

    print(Solution().removeBoxes([1, 1, 1]))  # 9
    print(Solution().removeBoxes([1]))  # 1

    # 自制用例
    print(Solution().removeBoxes([1, 2, 3, 4, 5, 6, 1, 6, 5, 4, 3, 2, 1]))  # 25
